import os,re
from modules import base
class Disk(base.Base):
    def __init__(self,host=None):
        self.command=r"lsblk | grep '^sd.' | awk '{print $1,$4}'"
        self.host=host
        self.res_list=[]
        super(Disk,self).__init__()

    def windows(self):
        res=re.findall('PHYSI.*\d+',os.popen('wmic DISKDRIVE get size,name').read())
        for disk_info in res:
            self.res_list.append({disk_info.split('  ')[0]:'%sG'%str(int(disk_info.split('  ')[1])/1000/1000/1000).split('.')[0]})
        base.data_dict['disk_info']=self.res_list
    def linux(self):
        if self.host==None: #如果=None,则证明是agent模式
            res=os.popen(self.command).read()
            self.parse(res)
            base.data_dict['disk_info']=self.res_list
            base.data_dict['disk_info']=self.res_list

        else: #传入了host_dict，证明是ssh模式
            res=self.ssh(self.host,self.command).decode(encoding='utf-8')
            self.parse(res)
            base.data_dict[self.host]['disk_info']=self.res_list
    def parse(self,data):
        new_data=data.split("\n")
        for i in new_data:
            tmp_dict={}
            clean=i.split(' ')
            if len(clean) >1:
                tmp_dict[clean[0]]=clean[1]
                self.res_list.append(tmp_dict)



